from IPython.parallel import Client
#c=Client(profile='rwth_physik')
#c=Client(profile='rwth_physik_tag')
#c=Client(profile='rwth_physik_prufen')
c=Client()    
dv=c[:]
lbv = c.load_balanced_view()
#dv=c[:10]
#dv.execute('from sys import path; path.append("/net/home/lxtsfs1/tpa/criger/Documents/temp_shared_data/trajectory_run_dir")')
dv.execute('from sys import path; path.append("/home/dcriger/Documents/2013/November/Toric_Code_Measurement/stochastic-master-equations/python")')
with dv.sync_imports():
    from numpy import linspace, complex128, append, zeros
    from utils import stoc_eval, traj_fid

def fid_append(var_tpl, var_names, gammas, gamma_names, pulse_f, input_dict):
    """
    returns a dictionary consisting of the input variables in var_tpl, with the names given in var_names, as well as a measurement fidelity obtained by evaluation until 
    """
    h_offset, h_peak, stdev, meas_t = \
        var_tpl[0], var_tpl[1]-var_tpl[0], var_tpl[2]*var_tpl[3], var_tpl[3]
    
    dep_dict = {name: var for name, var in zip(var_names, (h_offset, h_peak, stdev, meas_t))}
    
    dt = input_dict['dt']
    T = meas_t
    N = int(round(T / dt))
    tvec = linspace(0., T, N)
    
    epsilon = pulse_f(h_offset, h_peak, stdev, tvec)
    
    #Extreme Ugliness: add 2/chi to the time vector here, 
    #along with the requisite number of zeros to epsilon.

    T += 2 #Simulation time gets an extra 2/chi
    N = int(round(T / dt))
    tvec = linspace(0., T, N)

    zeros_len = len(tvec) - len(epsilon)
    epsilon = append(epsilon, zeros(zeros_len))

    input_dict['T'] = T
    input_dict['tvec'] = tvec
    input_dict['epsilon'] = epsilon
    input_dict['N'] = N
  
    gamma_d = gammas[0]
    gamma_r = gammas[1]
    Ld1 = input_dict['Ld1']; Ld2 = input_dict['Ld2']
    Ld3 = input_dict['Ld3']
    Lr1 = input_dict['Lr1']; Lr2 = input_dict['Lr2']
    Lr3 = input_dict['Lr3']
    kappa1 = input_dict['kappa1']; kappa2 = input_dict['kappa2']
    #Hard-coded Purcell decay factor 1/1600 
    D1 = (0.5 * gamma_d * Ld1 + 0.5 * gamma_d * Ld2 + 0.5 * gamma_d * Ld3 + 
            0 * (1./1600. * (kappa1 + kappa2) + gamma_r ) * (Lr1 + Lr2 + Lr3)).astype(complex128)

    input_dict.update({'D1': D1, 'eta': 0.5, 'tmeas_on': 0.,
                'tmeas_off': meas_t, 'epsilon': epsilon, 'N': N, 'T': meas_t, 'tvec':tvec})
    
    fid, ci = stoc_eval(traj_fid, (input_dict,), tol=10**-3, maxiter=10000)
    dep_dict.update({'fid':fid, 'confidence':ci})    
    return dep_dict

if __name__ == "__main__":
    import cPickle as pickle
    from itertools import product
    import time
    
    with dv.sync_imports():
        from utils import prep_dict_indep, current_timestamp
        from utils import gaussian_pulse as pulse_f
        #import utils as ut
        import numpy as np
    
    #To begin, prepare problem-independent dictionary
    input_dict = prep_dict_indep()
    '''
    h_offset_list = np.linspace(0.05,0.15,10)
    total_height_list = np.linspace(0.4,0.6,10)
    #h_peak_list = np.linspace(0.1,0.5,10)
    stdev_mul_list = np.linspace(0.1,1.0,10)
    meas_t_list = np.linspace(6.,15.,10)     
    '''
    h_offset_list = [0.1]#np.linspace(0.05,0.15,8)
    total_height_list = np.linspace(0.4,0.6,3)
    #h_peak_list = np.linspace(0.1,0.5,10)
    stdev_mul_list = [0.3]#np.linspace(0.1,0.5,5)
    meas_t_list = [10.] #np.linspace(6.,10.,2)     
    
    var_names = ['h_offset', 'h_peak', 'stdev', 'meas_t' ]

    gammas = [0., 0.] #Dephasing and relaxation rates
    gamma_names = ['gamma_d', 'gamma_r']

    dv.push({'fid_append':fid_append, 'var_names':var_names, 'gammas':gammas, 'gamma_names':gamma_names, 
                'pulse_f':pulse_f, 'input_dict':input_dict}) #TODO determine if loop vars need passing
    
    #lbv.block = True
    #dict_list = lbv.map_sync(lambda var_tpl: fid_append(var_tpl, var_names, gammas, gamma_names, pulse_f, input_dict), 
    #                    list(product(h_offset_list, total_height_list, stdev_mul_list, meas_t_list)))
    
    result_list = lbv.map_async(lambda var_tpl: fid_append(var_tpl, var_names, gammas, gamma_names, pulse_f, input_dict), 
                        list(product(h_offset_list, total_height_list, stdev_mul_list, meas_t_list)))
    
    dict_list = [{}]*len(result_list)
    """
    retrieved = [False] * len(dict_list)
    while True:
        all_done = True
        for ii, thing in enumerate(result_list):
            all_done = isinstance(thing, dict) and all_done
            if isinstance(thing, dict) and not retrieved[ii]:
                try:
                    test_keys = result_list[ii].keys()
                    dict_list[ii] = result_list[ii]
                    retrieved[ii] = True
                except:
                    pass
        time.sleep(1)

        if all_done:
            print 'done'
            break
    """
    result_list.wait()
    for idx in range(len(result_list)):
        try:
            dict_list[idx] = result_list[idx]
        except:
            pass

    with open("dict_list" + current_timestamp(),'w') as phil:
        pickle.dump(dict_list, phil)






















